package com.atlin.project.bizmq.bi;

import com.atlin.project.common.ErrorCode;
import com.atlin.project.exception.BusinessException;
import com.atlin.project.manager.AiManager;
import com.atlin.project.model.domain.chart.Chart;
import com.atlin.project.model.enums.chart.ChartStatusEnum;
import com.atlin.project.service.ChartService;
import com.atlin.project.service.ServiceUsageService;
import com.atlin.project.service.SysNotificationService;
import com.rabbitmq.client.Channel;
import io.github.briqt.spark4j.model.SparkSyncChatResponse;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.amqp.rabbit.annotation.RabbitListener;
import org.springframework.amqp.support.AmqpHeaders;
import org.springframework.messaging.handler.annotation.Header;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.io.IOException;

/**
 * @author atlin
 * @description 消费
 * @createDate 2024/11/17 17:17
 */
@Component
@Slf4j
public class BiMessageConsumer {

    @Resource
    private ChartService chartService;

    @Resource
    private ServiceUsageService serviceUsageService;

    @Resource
    private AiManager aiManager;

    @Resource
    private SysNotificationService sysNotificationService;
    


    @RabbitListener(queues = {BiMqConstants.BI_QUEUE_NAME}, ackMode = "MANUAL")
    @SneakyThrows
    public void receiveMessage(String message, Channel channel, @Header(AmqpHeaders.DELIVERY_TAG) long deliverTag) {
        log.info("receiveMessage message={}", message);
        if (StringUtils.isBlank(message)) {
            throwExceptionAndNackMessage(Long.parseLong(message), channel, deliverTag);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "消息为空");
        }
        genChart(Long.parseLong(message), channel, deliverTag);
    }

    /**
     * 图表生成
     */
    private void genChart(Long chartId, Channel channel, long deliverTag) {
        Chart chart = chartService.getById(chartId);
        if (chart == null) {
            throwExceptionAndNackMessage(chartId, channel, deliverTag);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "消息为空");
        }

        // 校验是否为自动重试
        boolean isRetry = chart.getStatus().equals(ChartStatusEnum.FAILURE.getValue());

        if (!isRetry) {
            if (!ChartStatusEnum.QUEUEING.getValue().equals(chart.getStatus())) {
                log.error("{}状态不为排队中", chart.getId());
                throwExceptionAndNackMessage(chartId, channel, deliverTag);
                return;
            }
            // 检查tokens 是否足够
            if (!checkLeftTokens(chart.getUserId())) {
                throwExceptionAndNackMessage(chartId, channel, deliverTag);
                chartService.handleChartUpdateError(chart.getId(), "剩余Tokens不足", ChartStatusEnum.UNDERBALANCE);
                return;
            }
        }

        //修改状态
        Chart updateChart = new Chart();
        updateChart.setId(chart.getId());
        updateChart.setStatus(ChartStatusEnum.GENERATING.getValue());
        boolean b = chartService.updateById(updateChart);
        if (!b) {
            throwExceptionAndNackMessage(chartId, channel, deliverTag);
            chartService.handleChartUpdateError(chart.getId(), "更新图表执行中状态失败");
            return;
        }

        String input = buildUserInput(chart, channel, deliverTag);
        if (StringUtils.isBlank(input)) return;

        // 调用AI
        SparkSyncChatResponse sparkSyncChatResponse = aiManager.sendMsgToSparkAndCount(true, input);
        if (sparkSyncChatResponse == null) {
            throwExceptionAndNackMessage(chartId, channel, deliverTag);
            chartService.handleChartUpdateError(chart.getId(), "AI生成失败");
            return;
        }
        Integer totalTokens = sparkSyncChatResponse.getTextUsage().getTotalTokens();

        String resultContent = sparkSyncChatResponse.getContent();
        String[] split = resultContent.split("【【【【【");
        if (split.length < 3) {
            throwExceptionAndNackMessage(chartId, channel, deliverTag);
            chartService.handleChartUpdateError(chart.getId(), "AI生成失败");
            return;
        }
        String genChart = split[1].trim();
        String genResult = split[2].trim();
        // option 校验 --- 失败重试
        boolean validator = chartService.chartOptionValidator(genChart);

        // 更新消耗的tokens
        long userId = isRetry ? 1L : chart.getUserId();
        boolean consumerCountRes = consumerTokenCount(userId, Long.valueOf(totalTokens));
        if (!consumerCountRes) {
            // 失败记录日志
            log.error("userId:{} consumerTokenCount,tokens:{} failure", userId, totalTokens);
        }

        // 存储得到的结果
        Chart updateChartResult = new Chart();
        updateChartResult.setId(chart.getId());
        updateChartResult.setGenChart(genChart);
        updateChartResult.setGenResult(genResult);
        updateChartResult.setStatus(ChartStatusEnum.SUCCESS.getValue());
        boolean updateResult = chartService.updateById(updateChartResult);
        if (!updateResult) {
            throwExceptionAndNackMessage(chartId, channel, deliverTag);
            chartService.handleChartUpdateError(chart.getId(), "更新图表成功状态失败");
        }
        //消息确认
        try {
            channel.basicAck(deliverTag, false);
        } catch (IOException e) {
            log.error("{} basicAck failed", chartId);
        }
        notifyUser(chart);
    }

    /**
     * 检查用户token是否足够
     *
     * @param userId 用户
     */
    private boolean checkLeftTokens(Long userId) {
        Long leftTokens = serviceUsageService.getLeftTokens(userId);
        return leftTokens >= 10;
    }

    /**
     * 统计用户此任务消耗的tokens
     *
     * @param userId 用户
     * @param tokens token
     */
    private boolean consumerTokenCount(Long userId, Long tokens) {
        log.info("userId:{} consumerTokenCount,tokens:{}", userId, tokens);
        return serviceUsageService.updateTokens(userId, tokens, true);
    }


    /**
     * 消息拒绝
     */
    private void throwExceptionAndNackMessage(Long chartId, Channel channel, long deliverTag) {
        try {
            channel.basicNack(deliverTag, false, false);
        } catch (IOException e) {
            log.error(e.getMessage());
        }
        log.error("basicNack {}", chartId);
    }

    /**
     * 通知客户端
     */
    private void notifyUser(Chart chart) {
        String notifyMessage = "您提交的智能分析任务 [ " + chart.getName() + " ] 已经完成，请前往我的图表查看。";
        sysNotificationService.notifyUser(chart.getUserId(), "智能分析", notifyMessage);
    }

    /**
     * 构建输入
     */
    private String buildUserInput(Chart chart, Channel channel, long deliverTag) {
        String goal = chart.getGoal();
        String chartData = null;
        try {
            chartData = chartService.getChartOriginalCsvData(chart.getId());
        } catch (Exception e) {
            throwExceptionAndNackMessage(chart.getId(), channel, deliverTag);
            chartService.handleChartUpdateError(chart.getId(), "图表数据查询失败");
            return "";
        }

        if (StringUtils.isBlank(chartData)) {
            throwExceptionAndNackMessage(chart.getId(), channel, deliverTag);
            chartService.handleChartUpdateError(chart.getId(), "图表数据为空");
            return "";
        }
        String chartType = chart.getChartType();

        // 拼接问题
        StringBuilder userInput = new StringBuilder();
        userInput.append("目标:").append(goal)
                .append("，请使用[").append(chartType).append("]图\n");
        userInput.append("数据:").append(chartData);
        return userInput.toString();
    }

}
